import base64
import json
import os
import sys
import time

import requests

import config


def fofa_search(search_files, search_type, current_path):
    ip_list = []
    unique_ips = set()
    search_files_name = search_files[:-4]
    search_files_names = "Fofa_search_" + search_files_name + ".txt"
    url_list_name = search_files_name + "_url_list.txt"
    url_ip_list = search_files_name + "ip_list.txt"
    pocbomber_ret = search_files_name + "_pocbomber_ret.txt"
    search_files_names = "{}\{}".format(current_path, search_files_names)
    url_list_name = "{}\{}".format(current_path, url_list_name)
    ip_list_name = "{}\{}".format(current_path, url_ip_list)
    with open(search_files, 'rt', encoding="utf-8") as files:
        for line in files:
            query = line.strip()
            query = '''{0}="{1}"'''.format(search_type, query)
            query_base64 = base64.b64encode(query.encode('utf-8')).decode('utf-8')
            url = '''https://fofa.info/api/v1/search/all?email={0}&key={1}&qbase64={2}&size=1000&fields=host,ip,port,server,os,title,product,version,city,region,icp
            '''.format(email, key, query_base64)
            try:
                res = requests.get(url, timeout=10).text
                if "账号无效" in res:
                    message = "[-]error 账号无效,请替换API KEY!!!"
                    print(message)
                    return message
                # if "权限" in res:
                #     message = res
                #     print(message)
                #     return message
                data = json.loads(res)
                title = data['query'].split('=')[1].replace('"', '')
                results = data['results']
                title = "{}".format(title)
                if results:
                    print("\n[+]{0} \n".format(title))
                    with open(search_files_names, "a+", encoding='utf-8') as f_title:
                        f_title.write("\n\n\n[+]{0}\n".format(title))
                    for result in results:
                        domain_name = result[0]
                        if "cdn" in domain_name:
                            break
                        if "CDN" in domain_name:
                            break
                        if "static" in domain_name:
                            break
                        if "STATIC" in domain_name:
                            break
                        ip = result[1]
                        ip_list.append(ip)
                        port = result[2]
                        server = result[3]
                        os = result[4]
                        title = result[5]
                        product = result[6]
                        version = result[7]
                        city = result[8]
                        region = result[9]
                        icp = result[10]
                        if title and title != "404 – Not Found":
                            with open(url_list_name, "a+", encoding='utf-8') as url_list_file:
                                url_list_file.write(domain_name + '\n')

                        with open(search_files_names, "a+", encoding='utf-8') as f:
                            print(
                                "      域名为：{}, IP: {}, PORT: {}, server:{}, os:{}, title:{}".format(
                                    domain_name, ip, port, server, os, title, product, version, city, region, icp))
                            data_info = "      域名为：{}, IP: {}:{}, server:{}, os:{}, title:{}".format(
                                domain_name, ip, port, server, os, title, product, version, city, region, icp)
                            f.write("{} \n".format(data_info))
                else:
                    print("\n[-]{0} \n".format(title))
                    print('      因为网络原因查询中断或该资产没有信息，请师傅们手工验证!' + '\n')
                    with open(search_files_names, "a+", encoding='utf-8') as f:
                        f.write("\n[-]{}\n".format(title))
                        f.write("      因为网络原因查询中断或该资产没有信息，请师傅们手工验证!" + '\n')
            except Exception as e:
                print(e)
                url_list_name = None
                print("[-] error 网络连接失败, 请查看网络设置后重试！！！")
                return url_list_name, pocbomber_ret

        for ip in ip_list:
            # 将ip地址添加到集合中
            unique_ips.add(ip)
        # 将去重后的ip地址写入文件
        with open(ip_list_name, "a+", encoding='utf-8') as ip_list_file:
            for ip in unique_ips:
                ip_list_file.write(ip + '\n')

        print("[*]-----------------------------------本次搜索完毕-----------------------------------------")
        print("[*]文件保存位置为:{}".format(current_path))
        return url_list_name, pocbomber_ret


def fofa_search_banner():
    banner = '''
     ████████            ████           ████████ ██  ██          ████████                           
░██░░░░░            ░██░           ░██░░░░░ ░░  ░██         ██░░░░░░                            
░██        ██████  ██████  ██████  ░██       ██ ░██  █████ ░██         █████   ██████   ███████ 
░███████  ██░░░░██░░░██░  ░░░░░░██ ░███████ ░██ ░██ ██░░░██░█████████ ██░░░██ ░░░░░░██ ░░██░░░██
░██░░░░  ░██   ░██  ░██    ███████ ░██░░░░  ░██ ░██░███████░░░░░░░░██░██  ░░   ███████  ░██  ░██
░██      ░██   ░██  ░██   ██░░░░██ ░██      ░██ ░██░██░░░░        ░██░██   ██ ██░░░░██  ░██  ░██
░██      ░░██████   ░██  ░░████████░██      ░██ ███░░██████ ████████ ░░█████ ░░████████ ███  ░██
░░        ░░░░░░    ░░    ░░░░░░░░ ░░       ░░ ░░░  ░░░░░░ ░░░░░░░░   ░░░░░   ░░░░░░░░ ░░░   ░░ 

                                                                         _______________________                         
                                                                        | author: Spike.        |
                                                                        | email: lendme@163.com |
                                                                        |_______________________| 
    '''
    print(banner)


def start_POC_bomber(url_list_name, pocbomber_ret):
    exec_POC_bomber = "python pocbomber.py -f {0} -t 10 -o {1} ".format(url_list_name, pocbomber_ret)
    os.system(exec_POC_bomber)


if __name__ == '__main__':
    fofa_search_banner()
    time.sleep(1)
    # 配置fofa登录邮箱和key
    email = config.email
    key = config.key
    try:
        search_files = sys.argv[1]
        search_type = sys.argv[2]
        file_type = search_files[-3:]
        search_files_name = search_files[:-4]
        current_path = os.getcwd()
        if os.path.exists(search_files_name):
            os.remove("{}\{}".format(current_path, search_files_name))
        os.mkdir(search_files_name)

        current_path = "{}\{}".format(current_path, search_files_name)
        if file_type == "txt":
            print("[*]-----------------------------------正在进行fofa匹配-----------------------------------")
            url_list_name, pocbomber_ret = fofa_search(search_files, search_type, current_path)
            # if url_list_name:
            #     if "账号无效" not in url_list_name:
            #         pocbomber_ret_path = "{}\{}".format(current_path, pocbomber_ret)
            #         start_POC_bomber(url_list_name, pocbomber_ret_path)
        else:
            print("[-]  文件读取错误，请使用txt文件！")
            print("[*]  usage: python3 fofafilescan.py  您需要查找的文件.txt   查询的类型域名: domain, 网站名称: title")
    except Exception as e:
        print(e)
    sys.exit()
